#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ChromaDB数据查看脚本
用于验证conversation_create是否成功创建数据
"""

import chromadb
import yaml
import os
from datetime import datetime

def load_config():
    """加载数据库配置"""
    config_path = './config/database.yaml'
    if os.path.exists(config_path):
        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    return {}

def view_chromadb_data():
    """查看ChromaDB中的数据"""
    print("=" * 60)
    print("📊 ChromaDB 数据查看器")
    print("=" * 60)
    
    # 加载配置
    config = load_config()
    chromadb_config = config.get('chromadb', {})
    persist_directory = chromadb_config.get('persist_directory', './data/chroma_db')
    
    print(f"📁 数据库路径: {persist_directory}")
    print(f"📁 绝对路径: {os.path.abspath(persist_directory)}")
    
    try:
        # 连接ChromaDB
        client = chromadb.PersistentClient(path=persist_directory)
        
        # 列出所有集合
        collections = client.list_collections()
        print(f"\n📚 发现 {len(collections)} 个集合:")
        for collection in collections:
            print(f"   - {collection.name}")
        
        # 查看conversations集合
        if any(c.name == 'conversations' for c in collections):
            print("\n🔍 查看 'conversations' 集合:")
            collection = client.get_collection('conversations')
            count = collection.count()
            print(f"   总记录数: {count}")
            
            if count > 0:
                # 获取最近的10条记录
                results = collection.get(
                    limit=min(10, count),
                    include=['documents', 'metadatas']
                )
                
                print("\n📝 最近的记录:")
                for i, (doc, meta) in enumerate(zip(
                    results['documents'], 
                    results['metadatas']
                )):
                    conv_id = meta.get('conversation_id', 'N/A')
                    msg_id = meta.get('message_id', 'N/A')
                    timestamp = meta.get('timestamp', 'N/A')
                    role = meta.get('role', 'N/A')
                    
                    print(f"\n   {i+1}. 记录编号: {i+1}")
                    print(f"      对话ID: {conv_id}")
                    print(f"      消息ID: {msg_id}")
                    print(f"      角色: {role}")
                    print(f"      时间: {timestamp}")
                    print(f"      内容: {doc[:100]}{'...' if len(doc) > 100 else ''}")
            else:
                print("   ⚠️ 集合为空")
        else:
            print("\n❌ 未找到 'conversations' 集合")
        
        # 查看knowledge_base集合（如果存在）
        if any(c.name == 'knowledge_base' for c in collections):
            print("\n🔍 查看 'knowledge_base' 集合:")
            collection = client.get_collection('knowledge_base')
            count = collection.count()
            print(f"   总记录数: {count}")
            
            if count > 0:
                results = collection.get(
                    limit=min(5, count),
                    include=['documents', 'metadatas']
                )
                
                print("\n📝 最近的记录:")
                for i, (doc, meta) in enumerate(zip(
                    results['documents'], 
                    results['metadatas']
                )):
                    print(f"\n   {i+1}. 记录编号: {i+1}")
                    print(f"      元数据: {meta}")
                    print(f"      内容: {doc[:100]}{'...' if len(doc) > 100 else ''}")
        
    except Exception as e:
        print(f"❌ 连接ChromaDB失败: {e}")
        return False
    
    print("\n" + "=" * 60)
    print("✅ 数据查看完成")
    print("=" * 60)
    return True

if __name__ == "__main__":
    view_chromadb_data()